package com.newbee.aop;

import com.alibaba.fastjson.JSON;
import com.google.common.collect.ImmutableMap;
import com.newbee.builder.TraceNodeBuilder;
import com.newbee.cache.TraceNodeCache;
import com.newbee.enums.TraceNodeEnum;
import com.newbee.model.TraceContext;
import com.newbee.model.TraceNode;
import com.newbee.model.TraceNodeStack;
import com.newbee.model.common.CommonResponse;
import com.newbee.model.common.BaseResponse;
import com.newbee.utils.SensitiveUtil;
import com.newbee.utils.TraceContextHolder;
import com.newbee.utils.TraceNodeStackHolder;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.exception.ExceptionUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.lang.reflect.Method;
import java.util.Date;
import java.util.Objects;

@Aspect
@Component
@Slf4j
public class TraceAspect {

    @Autowired
    private TraceNodeBuilder traceNodeBuilder;

    @Autowired
    private TraceNodeCache traceNodeCache;

    @Resource(name = "traceContextShareConfig")
    private ImmutableMap<String, Boolean> traceContextShareConfig;

    @Resource(name = "traceEnable")
    private boolean traceEnable;

    @Pointcut("(execution(* com.newbee.service..*(..)) || execution(* com.newbee.controller..*(..)))")
    public void tracePointcut() {
    }

    @Around("tracePointcut()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        if (hasSkipTraceAnnotation(joinPoint)) {
            return joinPoint.proceed();
        }
        if (!traceEnable) {
            return joinPoint.proceed();
        }

        TraceContext context;
        TraceNode currentNode = null;
        TraceNodeStack stack = null;
        try {
            // 1. 获取线程上下文环境, 如果第一次则初始化
            context = TraceContextHolder.getOrInit(true);

            // 2.初始化当前节点
            currentNode = traceNodeBuilder.buildTraceNode(context, joinPoint);

            // 3. 获取节点栈, 如果第一次则初始化
            stack = TraceNodeStackHolder.getOrInit(currentNode);

            // 4. 把当前节点入栈
            stack.pushNodeAndRelatedChildNodes(currentNode);
        } catch (Throwable e) {
            log.error("tracePointcut异常", e);
        }

        long startTime = System.currentTimeMillis();
        // 5. 执行业务方法
        try {
            Object result = joinPoint.proceed();

            // 6. 更新成功状态（补充耗时日志，便于问题排查）
            if (Objects.nonNull(currentNode)) {
                currentNode.setStatus(TraceNodeEnum.SUCCESS);
                currentNode.setResponseParam(SensitiveUtil.desensitize(result));
                currentNode.setCostTime(System.currentTimeMillis() - startTime);
                currentNode.setUpdateTime(new Date());
                if (currentNode.getCostTime() > currentNode.getWarnThreshold()) {
                    log.warn("节点 {} 执行耗时过长：{} ms（阈值：{} ms）",
                            currentNode.getMethodName(), currentNode.getCostTime(), currentNode.getWarnThreshold());
                }

                // 需要返回traceId的情况
                // 当前节点是根节点, 按照约定, 返回值需要继承BaseResponse
                // 配置化
                if (Objects.nonNull(stack) && currentNode == stack.getRoot() && traceContextShareConfig.getOrDefault(currentNode.getMethodPath(), false)) {
                    if (result instanceof CommonResponse) {
                        Object data = ((CommonResponse<?>) result).getData();
                        if (data instanceof BaseResponse) {
                            BaseResponse baseResponse = (BaseResponse) data;
                            TraceContext traceContext = TraceContextHolder.getOrInit(false);
                            if (Objects.nonNull(traceContext)) {
                                baseResponse.setTraceId(traceContext.getTraceId());
                            }
                        }
                    }
                }
            }

            return result;
        } catch (Throwable e) {
            // 7. 更新失败状态（确保异常信息完整）
            if (Objects.nonNull(currentNode)) {
                currentNode.setStatus(TraceNodeEnum.FAIL);
                currentNode.setCostTime(System.currentTimeMillis() - startTime);
                currentNode.setUpdateTime(new Date());
                currentNode.setExceptionStack(ExceptionUtils.getStackTrace(e));

                log.error("节点 {} 执行失败（父节点：{}）错误信息: {}",
                        currentNode.getMethodName(),
                        null == stack ? null : stack.getParent(),
                        currentNode.getExceptionStack()
                );
            }

            throw e;
        } finally {
            // 8. 弹出当前节点
            if (Objects.nonNull(stack)) {
                TraceNode popNode = stack.popNode();
                if (Objects.nonNull(popNode)) {
                    log.info("{} 出栈", popNode.getMethodName());
                    // 9. 如果弹出的节点就是根节点, 说明链路已结束（null用于兼容异常情况）
                    if (popNode == stack.getRoot()) {
                        traceNodeCache.addNode(popNode);
                        TraceContextHolder.clear();
                        TraceNodeStackHolder.clear();
                        log.info("根节点执行完成, 链路已缓存: {}", JSON.toJSONString(currentNode));
                    }
                }
            } else {
                TraceContextHolder.clear();
                TraceNodeStackHolder.clear();
                log.warn("根节点执行异常, 已清理上下文: {}", JSON.toJSONString(currentNode));
            }
        }
    }

    /**
     * 判断当前方法是否标记了@SkipTrace注解
     */
    private boolean hasSkipTraceAnnotation(ProceedingJoinPoint joinPoint) {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        SkipTrace skipTrace = method.getAnnotation(SkipTrace.class);

        if (skipTrace != null) {
            log.info("方法[{}]标记@SkipTrace，跳过追踪（原因：{}）",
                    method.getName(), skipTrace.reason());
            return true;
        }
        return false;
    }
}